
#include <stdint.h>
#include <x264.h>

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <inttypes.h>
#include <unistd.h>
#include <sys/time.h>

uint64_t now(void)
{
    uint64_t ts = 0;

    struct timeval tv;
    memset(&tv, 0, sizeof(tv));
    gettimeofday(&tv, NULL);
    ts = (uint64_t)tv.tv_sec * 1000000 + tv.tv_usec;

    return ts;
}

typedef struct avc_encoder
{
    x264_param_t x264_param;
    x264_t* x264;

    x264_picture_t in_image;
    x264_picture_t out_frame;

    int width;
    int height;
    int bitrate;
    int fps;

    int gop_size;
    int gop_frame_count;

    void (*on_frame_data)(void* frame, unsigned len, void* userdata);
    void *userdata;
}avc_encoder_t;

typedef struct nv12_image
{
    int width;
    int height;
    unsigned char* data;
    int size;
    unsigned char* y_ptr;
    int y_stride;
    unsigned char* uv_ptr;
    int uv_stride;
}nv12_image_t;

nv12_image_t* nv12_image_alloc(int width, int height)
{
    int size = (width * height * 3) >> 1;
    nv12_image_t* image;

    image = (nv12_image_t*)malloc(sizeof(*image) + size);

    image->width = width;
    image->height = height;
    image->data = (unsigned char*)&image[1];
    image->size = size;

    image->y_ptr = image->data;
    image->y_stride = width;
    image->uv_ptr = image->y_ptr + width * height;
    image->uv_stride = width;

    return image;
}

void nv12_image_free(nv12_image_t* image)
{
    free(image);
    return;
}

#if 0
void fill_nv12_image(nv12_image_t* image, int index)
{
    int x, y;
    unsigned char* y_ptr;
    unsigned char* uv_ptr;
    int uv_width, uv_height;

    /* Y */
    y_ptr = image->y_ptr;
    for (y = 0; y < image->height; ++y)
    {
        for (x = 0; x < image->width; ++x)
        {
            y_ptr[x] = (unsigned char)(x + y + index * 3);
        }
        y_ptr += image->y_stride;
    }

    /* UV */
    uv_ptr = image->uv_ptr;
    uv_width = image->width;
    uv_height = image->height>>1;
    for (y = 0; y < uv_height; y++)
    {
        for (x = 0; x < uv_width; x += 2)
        {
			uv_ptr[x] = (unsigned char)(128 + y + index * 2);
			uv_ptr[x+1] = (unsigned char)(64 + x + index * 5);
        }
        uv_ptr += image->uv_stride;
    }

    return;
}

#else

void fill_nv12_image(nv12_image_t* image, int index)
{
    int i, j;

    for (j = 0; j < image->height / 2; j++)
    {
        unsigned char *py = image->y_ptr + 2 * j * image->y_stride;
        unsigned char *puv = image->uv_ptr + j * image->uv_stride;

        for (i = 0; i < image->width / 2; i++)
        {
            int z = (((i + index) >> 4) ^ ((j + index) >> 4)) & 15;
            py[0] = py[1] = py[image->y_stride] = py[image->y_stride + 1] = 0x80 + z * 0x8;
            puv[0] = 0x00 + z * 0x10;
            puv[1] = 0x80 + z * 0x30;

            py += 2;
            puv += 2;
        }
    }

    return;
}

#endif

avc_encoder_t* avc_encoder_new
(
    int width, int height, int fps, int bitrate,
    void (*on_frame_data)(void* frame, unsigned len, void* userdata), void *userdata
)
{
    avc_encoder_t* encoder;

    encoder = malloc(sizeof(*encoder));
    memset(encoder, 0, sizeof(*encoder));
    encoder->width = width;
    encoder->height = height;
    encoder->fps = fps;
    encoder->bitrate = bitrate;

    encoder->on_frame_data = on_frame_data;
    encoder->userdata = userdata;

    encoder->gop_size = 25;
    encoder->gop_frame_count = 0;

    return encoder;
}

int avc_encoder_init(avc_encoder_t* encoder)
{
    int ret;

    x264_param_default(&encoder->x264_param);

    //x264_param_default_preset(&encoder->x264_param, "veryfast", "animation");

    encoder->x264_param.i_width = encoder->width;
    encoder->x264_param.i_height = encoder->height;
    encoder->x264_param.i_csp = X264_CSP_NV12;

    encoder->x264_param.rc.i_rc_method = X264_RC_ABR;

    /* QP 确定画面细节保留程度，取值越小，细节越丰富 */
    encoder->x264_param.rc.i_qp_min = 18;
    encoder->x264_param.rc.i_qp_max = 30;
    encoder->x264_param.rc.i_qp_step = 3;

    encoder->x264_param.rc.i_bitrate = encoder->bitrate;

    encoder->x264_param.b_repeat_headers = 1;
    encoder->x264_param.b_annexb = 1;

    encoder->x264_param.i_fps_num = encoder->fps;
    encoder->x264_param.i_fps_den = 1;

    encoder->x264_param.i_timebase_num  = 1;
    encoder->x264_param.i_timebase_den  = encoder->fps;

    encoder->x264_param.i_frame_reference = 1;
    encoder->x264_param.i_dpb_size = 1;
    encoder->x264_param.b_intra_refresh = 0; /*  need i_frame_reference and  i_dpb_size to be both 1 to enable intra_refresh option */

    encoder->x264_param.i_keyint_max = encoder->fps;

    encoder->x264_param.i_bframe = 0;

    encoder->x264_param.b_deblocking_filter = 1;

    /* 设置 NALU 尺寸 */
    /* encoder->x264_param.i_slice_max_size = 1024; */

    /* more options */

    ret = x264_param_apply_profile(&encoder->x264_param, "high");
    if (ret != 0)
    {
        printf("x264_encoder_init, x264_param_apply_profile() failed, ret: 0x%X\n", ret);
        return -1;
    }

    encoder->x264 = x264_encoder_open(&encoder->x264_param);
    if (NULL == encoder->x264)
    {
        printf("x264_encoder_init, x264_encoder_open() failed\n");
        return -1;
    }

    return 0;
}

int avc_encoder_encode(avc_encoder_t* encoder, nv12_image_t* image, int pts, int keyframe)
{
    int ret;
    int i;

    x264_nal_t *p_nals;
    int nals_count;

    if (NULL != image)
    {
        encoder->gop_frame_count++;
        x264_picture_init(&encoder->in_image);

        if (keyframe)
        {
            encoder->in_image.i_type = X264_TYPE_IDR;
        }
        else
        {
            encoder->in_image.i_type = X264_TYPE_AUTO;
          #if 0
            if (encoder->gop_frame_count >= encoder->gop_size)
            {
                encoder->gop_frame_count = 0;
                x264_encoder_intra_refresh(encoder->x264);
            }
          #endif
        }

        encoder->in_image.i_pts = pts;

        encoder->in_image.img.i_csp = X264_CSP_NV12;
        encoder->in_image.img.i_plane = 2;
        encoder->in_image.img.i_stride[0] = image->y_stride;
        encoder->in_image.img.plane[0] = image->y_ptr;
        encoder->in_image.img.i_stride[1] = image->uv_stride;
        encoder->in_image.img.plane[1] = image->uv_ptr;

        p_nals = NULL;
        nals_count = 0;
        ret = x264_encoder_encode(encoder->x264, &p_nals, &nals_count, &encoder->in_image, &encoder->out_frame);
        if (ret < 0)
        {
            printf("avc_encoder_encode, x264_encoder_encode() failed, ret: 0x%X\n", ret);
            return -1;
        }

        for (i = 0; i < nals_count; ++i)
        {
            encoder->on_frame_data(p_nals[i].p_payload, (p_nals[i].i_payload), encoder->userdata);
        }
    }
    else
    {
        while (x264_encoder_delayed_frames(encoder->x264))
        {
            p_nals = NULL;
            nals_count = 0;
            ret = x264_encoder_encode(encoder->x264, &p_nals, &nals_count, NULL, &encoder->out_frame);
            if (ret < 0)
            {
                printf("avc_encoder_encode, delay, x264_encoder_encode() failed\n");
                break;
            }

            for (i = 0; i < nals_count; ++i)
            {
                encoder->on_frame_data(p_nals[i].p_payload, (p_nals[i].i_payload), encoder->userdata);
            }
        }
    }

    return 0;
}

void avc_encoder_destroy(avc_encoder_t* encoder)
{
    if (NULL != encoder->x264)
    {
        x264_encoder_close(encoder->x264);
    }
    free(encoder);

    return;
}

/********************************************************************************/

void on_frame_data(void* frame, unsigned len, void* userdata)
{
    FILE* fp = (FILE*)userdata;
    fwrite(&len, sizeof(len), 1, fp);
    fwrite(frame, len, 1, fp);

    return;
}

int main(int argc, char *argv[])
{
    int width = 640;
    int height = 480;
    int count = 100;
    int fps = 30;
    int bitrate = 128;
    const char* out_file = "x264.h264";

    FILE* fp;

    avc_encoder_t* encoder;
    nv12_image_t *image;
    int i;

    uint64_t ts1, ts2;

    if (argc > 3)
    {
        width = atoi(argv[1]);
        height = atoi(argv[2]);
    }
    if (argc > 4)
    {
        count = atoi(argv[3]);
    }
    if (argc > 5)
    {
        fps = atoi(argv[4]);
    }
    if (argc > 6)
    {
        bitrate = atoi(argv[5]);
    }

    fp = fopen(out_file, "wb");

    encoder = avc_encoder_new(width, height, fps, bitrate, on_frame_data, fp);

    avc_encoder_init(encoder);

    image = nv12_image_alloc(width, height);

    ts1 = now();
    for (i = 0; i < count; ++i)
    {
        fill_nv12_image(image, i);
        /* avc_encoder_encode(encoder, image, i, ((i%30) == 0)); */
        avc_encoder_encode(encoder, image, i, 0);
    }

    avc_encoder_encode(encoder, NULL, 0, 0);

    ts2 = now();

    printf("%d frame encoded, time elapsed: %"PRIu64"us, fps: %.3f\n", count, (ts2-ts1), (count / (1.0 * (ts2-ts1) / 1000000)));

    nv12_image_free(image);
    avc_encoder_destroy(encoder);

    fclose(fp);

	return 0;
}
